#include <iostream>
#include "distance.hpp"
#include <vector>
#include <cstring>
#include <unordered_set>
#include <unordered_map>
#include <queue>

namespace MultipleHNSW {

    class BruteForceSearch {
    public:
        BruteForceSearch(int vec_dim, int max_element) : vec_dim_(vec_dim), max_element_(max_element) {
        }

        BruteForceSearch() {
        }

        void insert(float *item);

        std::vector<std::unordered_set<int>> search_knn(float *query, int k, int n_query);

        std::unordered_set<int> search_knn(float *query, int k);

        ~BruteForceSearch() {
            data_.clear();
        }

    private:

        int max_element_;
        std::vector<float *> data_;
        int vec_dim_;
    };

    void BruteForceSearch::insert(float *item) {
        for (int i = 0; i < max_element_; i++) {
            data_.push_back(item + i * vec_dim_);
        }
    }

    std::unordered_set<int> BruteForceSearch::search_knn(float *query, int k) {
        std::unordered_set<int> result;
        int max_element = data_.size();

        std::priority_queue<distance_cmp, std::vector<distance_cmp>, distance_min_heap> queue;

        for (int i = 0; i < max_element; ++i) {
            double dis = l2distance(query, data_[i], vec_dim_);
            distance_cmp distance_item(dis, i, i);
            queue.push(distance_item);
        }
        for (int i = 0; i < k; ++i) {
            distance_cmp p = queue.top();
            queue.pop();
            result.insert(p.label_);
        }
        return result;
    }

    std::vector<std::unordered_set<int>> BruteForceSearch::search_knn(float *query, int k, int n_query) {
        std::vector<std::unordered_set<int>> result;
        for (int i = 0; i < n_query; i++) {
            std::unordered_set<int> tmp_gnd = search_knn(query + i * vec_dim_, k);
            result.push_back(tmp_gnd);
        }
        return result;
    }

}